#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time : 2023/1/3 19:11
# @Author : HuangBenHao
import math

from py2neo import Graph
import random
from pprint import pprint


def search_entity(entity, limit, graph):
    sql = "match (n)-[r]-(e) where n.name=" + "'" + entity + "'" + " return n,r,e"
    data = graph.run(sql).data()

    if not data:
        return []

    result = [create_json(item['n'], item['r'], item['e']) for item in data]
    result = average_type(result, int(limit))

    return result


def average_type(data, limit):
    result = []
    label_to_item = {}
    for item in data:
        end_type = item['p']['end']['labels'][0]
        if end_type not in label_to_item.keys():
            label_to_item[end_type] = []
        label_to_item[end_type] = label_to_item[end_type] + [item]
    temp = [{'key': key, 'count': len(label_to_item[key])} for key in label_to_item.keys()]
    temp.sort(key=lambda x: x['count'])
    average_count = math.floor(limit / len(temp)) + 1
    count = 0
    for item in temp:
        count = count + 1
        cur_data = label_to_item[item['key']]
        if count == len(temp):
            random.shuffle(cur_data)
            result = result + cur_data[:limit-len(result)]
            break
        if len(cur_data) <= average_count:
            result = result + cur_data
        else:
            random.shuffle(cur_data)
            result = result + cur_data[:average_count]
    return result


def create_json(n, r, e):
    result = {"p": {}}
    start = {}
    end = {}
    segment = {}
    if n:
        start['identity'] = n.identity
        start['labels'] = [str(n.labels).replace(':', '')]
        start['properties'] = {}
        for item in dict(n):
            start['properties'][item] = n[item]

    if e:
        end['identity'] = e.identity
        end['labels'] = [str(e.labels).replace(':', '')]
        end['properties'] = {}
        for item in dict(e):
            end['properties'][item] = e[item]

    if r:
        segment['relationship'] = {}
        segment['relationship']['identity'] = r.identity
        segment['relationship']['start'] = start['identity']
        segment['relationship']['end'] = end['identity']
        segment['relationship']['type'] = e['name']
        segment['start'] = start
        segment['end'] = end
        segment['relationship']['properties'] = {}
        for item in dict(r):
            segment['relationship']['properties'][item] = r[item]

    result['p']['start'] = start
    result['p']['end'] = end
    result['p']['segments'] = [segment]

    return result


if __name__ == '__main__':
    graph = Graph("http://localhost:7474", password="bigben123")
    entity = "糖尿病"
    limit = 50
    search_entity(entity, limit, graph)
    # print([1]+[2])
